import os
from torch.utils.data import Dataset
import cv2
import numpy as np
import torch
from ..misc.augment import PhotometricDistort, RandomSampleCrop


class DetectSingle(Dataset):
    def __init__(self, db_dir, shape=None):
        self.db_dir = db_dir
        self.data = []
        self.shape = shape
        for file in os.listdir(db_dir):
            if file.endswith(".png"):
                self.data.append(os.path.join(file))

    def __len__(self):
        return len(self.data)

    def _reader(self, path):
        image = cv2.imread(os.path.join(self.db_dir, path), cv2.IMREAD_GRAYSCALE)
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        return image.astype(np.float32), os.path.basename(path).split(".")[0]

    def __getitem__(self, index):
        img, name = self._reader(self.data[index])
        img = img / 127.5 - 1
        if self.shape:
            raw_shape = img.shape
            img = cv2.resize(img, (self.shape, self.shape))
        return {
            "images": img,
            "name": name,
            "shape": raw_shape if self.shape else img.shape,
        }


class ConditionSingle(Dataset):
    def __init__(self, db_dir, contrast=0.3):
        self.db_dir = db_dir
        self.contrast = contrast
        self.data = []
        for file in os.listdir(db_dir):
            if file.endswith(".png"):
                self.data.append(os.path.join(file))

    def __len__(self):
        return len(self.data)

    def _reader(self, path):
        image = cv2.imread(os.path.join(self.db_dir, path), cv2.IMREAD_GRAYSCALE)
        image = np.ones_like(image) * 255 / 2
        n = np.random.randint(5, 20)
        for _ in range(n):
            x = np.random.randint(0, image.shape[1] - 1)
            y = np.random.randint(0, image.shape[0] - 1)
            r = np.random.randint(5, 20)
            cv2.circle(image, (x, y), r, int(self.contrast * 255), -1)
        return image.astype(np.float32)

    def __getitem__(self, index):
        img = self._reader(self.data[index])[:, :, None]
        img = img / 127.5 - 1
        return {"image": img.repeat(3, 2)}


class TransferPairs(Dataset):
    def __init__(self, db_dir_A, db_dir_B):
        self.db_dir_A = db_dir_A
        self.db_dir_B = db_dir_B
        self.data_A = []
        self.data_B = []
        for file in os.listdir(db_dir_A):
            if file.endswith(".png"):
                self.data_A.append(os.path.join(file))
        for file in os.listdir(db_dir_B):
            if file.endswith(".png"):
                self.data_B.append(os.path.join(file))

    def __len__(self):
        return min(len(self.data_A), len(self.data_B))

    def _reader(self, path):
        image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        image = image / 127.5 - 1
        return image.astype(np.float32)

    def __getitem__(self, index):
        path_A = os.path.join(self.db_dir_A, self.data_A[index])
        path_B = os.path.join(self.db_dir_B, self.data_B[index])
        img_A = torch.tensor(self._reader(path_A))
        img_B = torch.tensor(self._reader(path_B))
        return {
            "A": img_A.unsqueeze(0).repeat(3, 1, 1),
            "B": img_B.unsqueeze(0).repeat(3, 1, 1),
        }


class DetectPairsWithTransfer(Dataset):
    def __init__(self, db_dir_A, db_dir_B, use_light=False, use_crop=False):
        self.db_dir_A = db_dir_A
        self.db_dir_B = db_dir_B
        self.data_A = []
        self.data_B = []
        self.light_trans = self.crop_trans = None
        if use_light:
            self.light_trans = PhotometricDistort()
        if use_crop:
            self.crop_trans = RandomSampleCrop()
        for file in os.listdir(db_dir_A):
            if file.endswith(".png"):
                self.data_A.append(os.path.join(file))
        for file in os.listdir(db_dir_B):
            if file.endswith(".png"):
                self.data_B.append(os.path.join(file))

    def __len__(self):
        return min(len(self.data_A), len(self.data_B))

    def _detect_reader(self, path):
        image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        with open(path.replace(".png", ".txt")) as f:
            bbox = []
            f.readline()
            for line in f:
                line = line.strip().split(",")
                bbox.append([int(x) for x in line])
        return image, np.array(bbox)

    def _image_reader(self, path):
        image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        return image

    def __getitem__(self, index):
        path_A = os.path.join(self.db_dir_A, self.data_A[index])
        path_B = os.path.join(self.db_dir_B, self.data_B[index])
        img_A, bbox = self._detect_reader(path_A)
        bbox[:, 2:] = bbox[:, 2:] + bbox[:, :2]
        if self.light_trans:
            img_A = self.light_trans(img_A)
        if self.crop_trans:
            img_A, bbox = self.crop_trans(img_A, bbox)
        img_A = (img_A / 127.5 - 1).astype(np.float32)
        img_B = self._image_reader(path_B)
        img_B = (img_B / 127.5 - 1).astype(np.float32)
        return {
            "trainA": img_A,
            "boxes": bbox,
            "trainB": img_B,
            "labels": np.ones(len(bbox), dtype=np.int64),
        }


class DetectWithBoxes(Dataset):
    def __init__(self, db_dir, shape=256):
        self.db_dir = db_dir
        self.data = []
        self.shape = shape
        for file in os.listdir(db_dir):
            if file.endswith(".png"):
                self.data.append(os.path.join(file))

    def __len__(self):
        return len(self.data)

    def _detect_reader(self, path):
        image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        with open(path.replace(".png", ".txt")) as f:
            bbox = []
            for line in f:
                line = line.strip().split(",")
                bbox.append([round(float(x)) for x in line])
        return image, np.array(bbox)

    def __getitem__(self, index):
        path = os.path.join(self.db_dir, self.data[index])
        img, bbox = self._detect_reader(path)
        img = (img / 127.5 - 1).astype(np.float32)
        bbox = bbox.astype(np.float32)
        bbox[:, [0, 2]] = bbox[:, [0, 2]] / img.shape[1] * self.shape
        bbox[:, [1, 3]] = bbox[:, [1, 3]] / img.shape[0] * self.shape
        img = cv2.resize(img, (self.shape, self.shape))
        return {
            "images": img,
            "boxes": bbox,
            "labels": np.ones(len(bbox), dtype=np.int64),
        }
